"""
@Time: 2021/1/24 下午 8:05
@Author: jinzhuan
@File: re_toolkit.py
@Desc: 
"""
import torch
from cognlp import *


torch.cuda.set_device(5)
tokenize_toolkit = TokenizeToolkit()
words = tokenize_toolkit.run('She would be the first foreign woman to die in the wave of kidnappings in Iraq.')

ner_toolkit = NerToolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/ner/trex/toolkit/model.pkl',
    vocabulary_path='../../../cognlp/data/ner/trex/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[5],
    max_seq_length=256)

ner_result = ner_toolkit.run(words)

re_toolkit = ReToolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/re/trex/toolkit/model.pt',
    vocabulary_path='../../../cognlp/data/re/trex/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[5],
    max_seq_length=256)

re_result = re_toolkit.run(words, ner_result)
print(re_result)
